"""
Jieun Oh

smr2.py -- Jieun's implementation of smr function for Music 422 HW4 #1(e)3

"""
import numpy as np
#import pylab
import window
import mdct
import psychoac as ps

def CalcSMRs(data, MDCTdata, MDCTscale, fs, sfBands):
    """Set SMR for each critical band in sfBands."""
    SMR = np.zeros(sfBands.nBands)
    N = data.size
    transform = abs(np.fft.rfft(window.HanningWindow(data),N)) #/N*4.4209        # not normalized
    transform = (transform**2) * 4 / ((N/2)**2) / (3./8.)

    x_axis = np.arange(0, N/2)*fs/N                      # for plot
    #pylab.plot(transform)
    #pylab.xlim(0, 1024)                               #
    #pylab.ylim(0, 1)                  




    # ------------- find peaks ----------------#
    peaks = np.array([])
    for i in range( np.cast['int'](N*0.5-3)):
        first = transform[i]
        second = transform[i+1]
        third = transform[i+2]
        if first<second and second>third and second>0.001: #0.001:
            peaks = np.append(peaks, i+1)
    
    fVec = np.arange(len(MDCTdata))*fs/N
    zVec = ps.Bark(fVec)
    


    # create one comprehensive masking curve
    alpha = 1.0
    ThreshDB = ps.Thresh(fVec)
    ThreshLin = ps.Intensity(ThreshDB)
    sumOfMasks = ThreshLin
 
    for i in range(len(peaks)):
        intensity_before = transform[peaks[i]-1]
        intensity_curr = transform[peaks[i]]
        intensity_after = transform[peaks[i]+1]
        intensity_sum = intensity_before + intensity_curr + intensity_after
        freq_before = (peaks[i]-1)*fs/N
        freq_curr = (peaks[i])*fs/N
        freq_after = (peaks[i]+1)*fs/N

        peakF = freq_before*intensity_before/intensity_sum + freq_curr*intensity_curr/intensity_sum + freq_after*intensity_after/intensity_sum  
        peakSPL = ps.SPL( intensity_before + intensity_curr + intensity_after  )
        #print ("peaks[", i, "] has SPL of ", peakSPL, " at frequency ", peakF)         
        m = ps.Masker(peakF, peakSPL)
        curve_i = m.vIntensityAtBark(zVec)              # NOT in dB
        
        #pylab.semilogx(x_axis, ps.SPL(curve_i), ':')    # for plot: masking by component i
        #pylab.hold(True)                                #

        curve_i = pow(curve_i, alpha)                 

        sumOfMasks = sumOfMasks + curve_i

    sumOfMasksDB = ps.SPL(pow(sumOfMasks, 1/alpha)) 
    maskSum = sumOfMasksDB


    #pylab.semilogx(x_axis, ps.Thresh(x_axis), '--')     # for plot: Threshold in Quiet
    #pylab.hold(True)                                    #

    #pylab.semilogx(x_axis, ps.SPL(transform[0:N/2]), '-r')   # for plot: original signal

    #pylab.semilogx(x_axis, sumOfMasksDB, '-b')          # for plot: final masked threshold
    #pylab.xlim(50, 20000)                               #
    #pylab.ylim(-10, 100)                                #
    #pylab.title('masking curves using alpha=1.00')      #
    #pylab.show()                                        #
    #pylab.hold(False)

    
    """ Comment out my code 

    for j in range(sfBands.nBands):          # each critical band
        mdctLine = sfBands.lowerLine[j]
        maxSMR = -999.0
        while mdctLine <= sfBands.upperLine[j] and mdctLine<len(MDCTdata):
           # print ("mdctline = ", mdctLine)
           ### curr = sumOfMasksDB[mdctLine]  - ps.SPL(transform[mdctLine])
	   ### curr =  ps.SPL(transform[mdctLine])- sumOfMasksDB[mdctLine] 
	   ### curr =  ps.SPL(MDCTdata[mdctLine])- sumOfMasksDB[mdctLine] 

           signalDB = ps.SPL(MDCTdata[mdctLine]*1.0/pow(2, MDCTscale))
	   curr =  signalDB- sumOfMasksDB[mdctLine] 


           if curr>maxSMR:
               maxSMR = curr
           mdctLine = mdctLine+1
        SMR[j] = maxSMR
    End of my code """ 


    """ Isaac's code """
    MDCTdata = abs(MDCTdata)
    MDCTdata = MDCTdata/(2**MDCTscale)
    
    # Find MDCT dB
    MDCTdB = ps.SPL(MDCTdata)
    
    # Check each band in sfBands and find signal-to-mask ratio
    for i in xrange(sfBands.nBands):
        if sfBands.upperLine[i] <= np.size(MDCTdata):
            xMaxb = MDCTdB[sfBands.lowerLine[i]:sfBands.upperLine[i]]
            Mb = max(maskSum[sfBands.lowerLine[i]:sfBands.upperLine[i]])
            SMR[i] = max(xMaxb - Mb)
    """ End of Isaac's code """


    #print SMR     # debug
    return SMR



# --------------------------Testing--------------------------------#

if __name__=="__main__":
    
    fs = 48000.0;
    N = 2048;                 
    t = np.arange(0,N)/fs; 

    f0 = 440;
    f1 = 550;
    f2 = 660;
    f3 = 880;
    f4 = 4400;
    f5 = 8800;

    a0 = 0.4;
    a1 = 0.2;
    a2 = 0.2;
    a3 = 0.09;
    a4 = 0.06;
    a5 = 0.05;

    data = (a0*np.cos(f0*2*np.pi*t) + a1*np.cos(f1*2*np.pi*t) + a2*np.cos(f2*2*np.pi*t) + a3*np.cos(f3*2*np.pi*t) + a4*np.cos(f4*2*np.pi*t) + a5*np.cos(f5*2*np.pi*t)); 
    dataW = window.HanningWindow(data)

    #pylab.plot(data)
    #pylab.show()

    mdctdata = mdct.MDCT(dataW, N/2, N/2)

    nLines = ps.AssignMDCTLinesFromFreqLimits(N, fs)
    sfBands = ps.ScaleFactorBands(nLines)

    SMR = CalcSMRs(data, mdctdata, 1, fs, sfBands)
    print SMR


    
